#include "stdafx.h"
#include "win_token.h"

#include <Wtsapi32.h>
#pragma comment(lib, "Wtsapi32.lib")

#include <TlHelp32.h>

_SSL_BEGIN
_SSWINOSL_BEGIN

bool token_enable_process_privilege(const tchar *privilege_name, bool enable)
{
	bool ret = false;
	HANDLE token = NULL;
	TOKEN_PRIVILEGES tkps = {0};
	tkps.PrivilegeCount = 1; // THE TOKEN_PRIVILEGES's MEMORY IS ONLY 1 PRIVILEGE

	if(!::OpenProcessToken(::GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES|TOKEN_QUERY, &token)
		|| token == NULL)
		goto finish;

	if(!::LookupPrivilegeValue(NULL, privilege_name, &tkps.Privileges[0].Luid))
		goto finish;

	tkps.Privileges[0].Attributes = enable ? SE_PRIVILEGE_ENABLED : SE_PRIVILEGE_REMOVED;

	if(!::AdjustTokenPrivileges(token, FALSE, &tkps, sizeof(tkps), NULL, NULL))
		goto finish;

	ret = true;

finish:
	if(token != NULL)
	{
		::CloseHandle(token);
		token = NULL;
	}

	return ret;
}

bool token_get_accout_names(tstring &name4process, tstring &name4session)
{
	name4process.clear();
	name4session.clear();

	unsigned long len = 0;
	if(::GetUserName(NULL, &len) || ::GetLastError() != ERROR_INSUFFICIENT_BUFFER)
		return false;

	std::auto_ptr<tchar> buffer(new tchar [len]);
	if(!::GetUserName(buffer.get(), &len))
		return false;

	name4process = buffer.get();

	tchar *name_ptr = NULL;
	if(!::WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE,
		WTS_CURRENT_SESSION,
		WTSUserName,
		&name_ptr,
		&len))
	{
		return false;
	}

	name4session = name_ptr;

	::WTSFreeMemory(name_ptr);
	name_ptr = NULL;

	return true;
}

HANDLE token_get_by_sessionid_and_accountname(unsigned long session_id, const tchar *account_name)
{
	if(account_name == NULL)
		return NULL;

	HANDLE snapshot_h = ::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if(snapshot_h == INVALID_HANDLE_VALUE)
		return NULL;

	HANDLE token = NULL;

	PROCESSENTRY32 pe32 = { sizeof(pe32) };
	if(::Process32First(snapshot_h, &pe32))
	{
		do 
		{
			unsigned long cur_session_id = 0;
			if(!::ProcessIdToSessionId(pe32.th32ProcessID, &cur_session_id))
				continue;
			if(cur_session_id != session_id)
				continue;
			HANDLE process_h = ::OpenProcess(PROCESS_ALL_ACCESS, TRUE, pe32.th32ProcessID);
			if(process_h == NULL)
				continue;

			HANDLE process_token_h = NULL;
			if(::OpenProcessToken(process_h,
				TOKEN_READ|TOKEN_DUPLICATE |TOKEN_QUERY|TOKEN_ASSIGN_PRIMARY,
				&process_token_h))
			{
				unsigned long len = 0;
				::GetTokenInformation(process_token_h, TokenUser, NULL, NULL, &len);
				if(len > 0)
				{
					std::auto_ptr<unsigned char> buffer(new unsigned char[len]);
					TOKEN_USER *token_info = (TOKEN_USER*)buffer.get();
					if(::GetTokenInformation(process_token_h, TokenUser, token_info,
						len, NULL))
					{
						SID_NAME_USE type = SidTypeUnknown;
						unsigned long username_len = 0, domain_len = 0;
						::LookupAccountSid(NULL, token_info->User.Sid, NULL, &username_len, NULL, &domain_len, &type);
						if(username_len > 0)
						{
							std::auto_ptr<tchar> user_name(new tchar[username_len]);
							std::auto_ptr<tchar> domain_name(new tchar[domain_len]);
							if(::LookupAccountSid(NULL, token_info->User.Sid, user_name.get(), &username_len,
								domain_name.get(), &domain_len, &type))
							{
								if(_tcscmp(account_name, user_name.get()) == 0)
								{
									token = process_token_h;
									process_token_h = NULL;
								}
							}
						}
					}
				}
			}
			::CloseHandle(process_token_h);
			::CloseHandle(process_h);
		} while (token == NULL && ::Process32Next(snapshot_h, &pe32));
	}

	::CloseHandle(snapshot_h);
	return token;
}

_SSWINOSL_END
_SSL_END

